import time
import json
from py2neo import Graph, Node, Relationship
from concurrent.futures import ThreadPoolExecutor
import threading

# 连接到Neo4j数据库
graph = Graph("http://localhost:7474/",auth=("neo4j", "123456"))


# 时间转换
def time_trans(now):
    timeArray = time.localtime(now)
    Time = time.strftime("%Y-%m-%d %H:%M:%S", timeArray)
    return Time


# 定义处理单行数据的函数
def process_data(data):

    # 提取src_ip、dst_ip和dst_port字段
    src_ip = data["src_ip"]
    dst_ip = data["dst_ip"]
    dst_port = data["dst_port"]
    in_packet_bytes = data["in_packet_bytes"]
    out_packet_bytes = data["out_packet_bytes"]
    conn_time = time_trans(float(data['__time__']))

    # 检查关系是否存在
    query = """
    MATCH (start:IP {ip: $src_ip})-[rel:CONNECTS_TO]->(end:IP {ip: $dst_ip})
    RETURN rel
    """

    # with data_lock:
    result = graph.evaluate(query, src_ip=src_ip, dst_ip=dst_ip)
    if result is None:
        # 创建或获取src_ip节点
        src_node = Node("IP", ip=src_ip)
        graph.merge(src_node, "IP", "ip")

        # 创建或获取dst_ip节点
        dst_node = Node("IP", ip=dst_ip)
        graph.merge(dst_node, "IP", "ip")

        # 创建关系及属性
        rel = Relationship(src_node, "CONNECTS_TO", dst_node, dst_port=dst_port, in_packet_bytes = in_packet_bytes, out_packet_bytes = out_packet_bytes, conn_time = conn_time)
        graph.create(rel)
    else:
        # 更新关系的times属性值
        query = """
        MATCH (start:IP {ip: $src_ip}), (end:IP {ip: $dst_ip})
        CREATE (start)-[:CONNECTS_TO {dst_port: $dst_port, in_packet_bytes: $in_packet_bytes, out_packet_bytes: $out_packet_bytes, conn_time: $conn_time}]->(end)
        """
        graph.run(query, src_ip=src_ip, dst_ip=dst_ip, dst_port=dst_port, in_packet_bytes=in_packet_bytes, out_packet_bytes=out_packet_bytes, conn_time=conn_time)


# 读取data.json文件
with open("C:/Users/yinwe/Desktop/data.json", "r") as file:
    lines = file.readlines()

# 加入线程锁
# data_lock = threading.Lock()

# 使用线程池处理数据
# with ThreadPoolExecutor() as executor:
for line in lines:
    data = json.loads(line)
    # executor.submit(process_data, data)
    process_data(data)